Argmax

沿指定轴查找最大 topk 个值的索引。当 topk=1 时，该算子等价于 ArgMax。

\[Y_i = \underset{k}{\operatorname{argmax}} (X_{slice_i})\]

其中 \(X_{slice_i}\) 是输入张量中沿指定轴的一个切片，函数返回该切片中最大值的索引 \(k\)。

输入：

  • input - 输入数据地址。

  • output - 输出索引的数据地址，数据类型通常为int32。

  • output_value - (可选) 输出值的数据地址。

  • in_shape - 输入张量的维度信息数组。

  • in_strides - 输入张量的步长信息数组。

  • out_strides - 输出张量的步长信息数组。

  • arg_elements - 用于存放候选值的临时工作空间地址。

  • index - 用于存放候选索引的临时工作空间地址。

  • topk - 需要查找的最大值的数量。设置为1以执行ArgMax操作。

  • out_value - 是否返回数值的标志。若为非0，则 output_value 必须提供有效地址。

  • input_shape_size - 输入张量的维度数 (即 in_shape 数组的长度)。

  • axis - 执行查找操作的轴。

  • core_mask - 核掩码。

输出：

  • output - 存储索引的输出张量。

  • output_value - 如果 return_values 为 true，则此处存储找到的值。

支持平台：

FT78NE MT7004

备注

  • FT78NE 支持fp32

  • MT7004 支持fp16, fp32

共享存储版本:

void fp_argmax_s(float *input, void *output, float *output_value, int32_t *in_shape, int *in_strides, int *out_strides, float *arg_elements, int *index, int topk, int out_value, int input_shape_size, int axis, int core_mask)
void hp_argmax_s(half *input, void *output, half *output_value, int32_t *in_shape, int *in_strides, int *out_strides, half *arg_elements, int *index, int topk, int out_value, int input_shape_size, int axis, int core_mask)

C调用示例：

 1//FT78NE示例
 2#include <stdio.h>
 3#include <argmax.h>
 4int main(int argc, char* argv[]) {
 5    float *input = (float *)0xA0000000;          // input在DDR空间
 6    int *output = (int *)0xB0000000;             // output indices
 7    float *output_value = (float *)0xC0000000;     // output values
 8    float *arg_elements = (float *)0xD0000000;   // temp workspace 1
 9    int *index = (int *)0xE0000000;               // temp workspace 2
10
11    int in_shape[] = {2, 3, 4};                  // input shape: (2, 3, 4)
12    int in_strides[] = {12, 4, 1};               // input strides for contiguous layout
13    int out_strides[] = {4, 1};                  // output strides, shape is (2, 4)
14    int input_shape_size = 3;
15
16    int axis = 1;                                // 沿第1轴操作
17    int topk = 1;                                // ArgMax
18    int out_value = 1;                           // 同时返回值
19    int core_mask = 0xff;
20
21    fp_argmax_s(input, output, output_value, in_shape, in_strides, out_strides, arg_elements, index, topk, out_value, input_shape_size, axis, core_mask);
22    return 0;
23}

私有存储版本:

void fp_argmax_p(float *input, void *output, float *output_value, int32_t *in_shape, int *in_strides, int *out_strides, float *arg_elements, int *index, int topk, int out_value, int input_shape_size, int axis)
void hp_argmax_p(half *input, void *output, half *output_value, int32_t *in_shape, int *in_strides, int *out_strides, half *arg_elements, int *index, int topk, int out_value, int input_shape_size, int axis)

C调用示例：

 1//FT78NE示例
 2#include <stdio.h>
 3#include <argmax.h>
 4int main(int argc, char* argv[]) {
 5    float *input = (float *)0x10001000;          // input在DDR空间
 6    int *output = (int *)0x10002000;             // output indices
 7    float *output_value = (float *)0x10003000;     // output values
 8    float *arg_elements = (float *)0x10004000;   // temp workspace 1
 9    int *index = (int *)0x10005000;               // temp workspace 2
10
11    int in_shape[] = {2, 3, 4};                  // input shape: (2, 3, 4)
12    int in_strides[] = {12, 4, 1};               // input strides for contiguous layout
13    int out_strides[] = {4, 1};                  // output strides, shape is (2, 4)
14    int input_shape_size = 3;
15
16    int axis = 1;                                // 沿第1轴操作
17    int topk = 1;                                // ArgMax
18    int out_value = 1;                           // 同时返回值
19
20    fp_argmax_p(input, output, output_value, in_shape, in_strides, out_strides,
21                arg_elements, index, topk, out_value, input_shape_size, axis);
22    return 0;
23}